#include "CThreadPool.h"
#include <exception>
#if !_WIN32
#include <sys/epoll.h>
#include <sys/eventfd.h>
#include <unistd.h>
#include <signal.h>
#include <sys/resource.h>
#endif

namespace TB
{
	std::atomic_flag CThreadPool::m_oIsManagerRunning = ATOMIC_FLAG_INIT;

	std::unordered_map<std::shared_ptr<CThreadPool>, CThreadPool::SThreadPoolInfo> CThreadPool::m_oThreadPools;

	std::unordered_map<std::thread::id, std::shared_ptr<std::thread>> CThreadPool::m_oThreads;

	CMPSCQueue<std::function<void()>> CThreadPool::m_oManagerTasks;

	std::shared_ptr<std::thread> CThreadPool::m_poManagerThread;

	CTimerMgr CThreadPool::m_oTimerMgr;

	std::vector<std::shared_ptr<CThreadPool>> CThreadPool::m_oToRemovedThreads;

	volatile UINT64 CThreadPool::m_qwCurTime = m_oTimerMgr.GetCurTime();

	std::list<std::shared_ptr<CThreadPool::SPollerContainer>> CThreadPool::m_oPollPerMilliSecond;
	
	volatile bool CThreadPool::m_bManagerThreadStoped = false;

	std::chrono::steady_clock::time_point CThreadPool::m_oStartTime = std::chrono::steady_clock::now();

	std::atomic_flag CThreadPool::m_oIsManagerTaskQueueNotEmpty = ATOMIC_FLAG_INIT;

#if _WIN32
	std::binary_semaphore CThreadPool::m_oManagerTaskSemaphore(0);
#else
	CObjectResMgr<CThreadPool::CPollEventCallbackObj> CThreadPool::m_oPollEventCallbackObjMgr([](UINT16) {return std::make_unique<CThreadPool::CPollEventCallbackObj>(); });
	CSpinLockFlag CThreadPool::m_oPollEventCallbackSpinLockFlag;
	int CThreadPool::m_nEpollFD = epoll_create(8192);
	int CThreadPool::m_nEventFD = eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK);
	epoll_event CThreadPool::m_aEpollEvents[MAX_EPOLL_EVENT_COUNT] = {};
#endif

	CThreadPool::CThreadPool(bool bIsDefault):
		m_bIsDefault(bIsDefault)
	{
		m_poSemaphore = std::make_unique<std::counting_semaphore<>>(0);
	}

	CThreadPool::~CThreadPool()
	{
	}

	std::shared_ptr<CThreadPool> CThreadPool::NewThreadPool()
	{
		StartManagerThread();
		auto poPoint = new CThreadPool(false);
		std::shared_ptr<CThreadPool> poThreadPool(poPoint);
		_PostManagerTask([poThreadPool]() 
		{
			m_oThreadPools.insert(std::make_pair(poThreadPool, SThreadPoolInfo()));
		});
		return poThreadPool;
	}

	void CThreadPool::DestroyThreadPool(const std::shared_ptr<CThreadPool>& poThreadPool)
	{
		if (poThreadPool->IsDefault())
			throw std::runtime_error("cannot destroy default thread pool");
		if (!poThreadPool->m_bToBeDestroyed)
		{
			poThreadPool->m_bToBeDestroyed = true;
			poThreadPool->m_poSemaphore->release(poThreadPool->m_poSemaphore->max());
		}
	}

	void CThreadPool::StartTimer(UINT64 qwMilliSeconds, const std::function<void()>& cbOnTimer)
	{
		CTimerMgr& roTimerMgr = m_oTimerMgr;
		_PostManagerTask([&roTimerMgr, qwMilliSeconds, cbOnTimer]()
		{
			roTimerMgr.StartTimer(qwMilliSeconds, 1, [cbOnTimer](STimerUID, UINT32, UINT32)
			{
				cbOnTimer();
			});
		});
	}

	CThreadPool::DManagerThreadRun CThreadPool::DetachManagerThread()
	{
		if (!m_oIsManagerRunning.test_and_set(std::memory_order_acquire))
			return &CThreadPool::_ManagerThreadRun;
		return nullptr;
	}

	void CThreadPool::StopManagerThread()
	{
		m_bManagerThreadStoped = true;
	}

	void CThreadPool::StartManagerThread()
	{
		m_bManagerThreadStoped = false;
		if (!m_oIsManagerRunning.test_and_set(std::memory_order_acquire))
		{
			m_bManagerThreadStoped = false;
			m_poManagerThread = std::make_shared<std::thread>(&CThreadPool::_ManagerThreadRun);
		}
	}

	void CThreadPool::JoinManagerThread()
	{
		if (m_poManagerThread && m_poManagerThread->joinable())
			m_poManagerThread->join();
	}

	void CThreadPool::_ManagerThreadRun()
	{
#if !_WIN32
		sigset_t signal_mask;
		sigemptyset(&signal_mask);
		sigaddset(&signal_mask, SIGPIPE);
		pthread_sigmask(SIG_BLOCK, &signal_mask, NULL);

		rlimit sLimit = {};
		if (0 == getrlimit(RLIMIT_NOFILE, &sLimit))
		{
			sLimit.rlim_cur = sLimit.rlim_max;
			setrlimit(RLIMIT_NOFILE, &sLimit);
		}

		UINT64 qwEventCallback = CreatePollEventCallback([](UINT32)
		{ 
			UINT64 v = 0; 
			read(CThreadPool::m_nEventFD, &v, sizeof(v));
			CThreadPool::m_oIsManagerTaskQueueNotEmpty.clear(std::memory_order_release);
		});
		if (!AddFD(m_nEventFD, EPOLLIN | EPOLLERR, qwEventCallback))
			throw std::runtime_error("add eventfd failed");
#endif
		m_oTimerMgr.SetMaxCountPerPoll(10);
		const STimerUID& sPollTimerUID = m_oTimerMgr.StartTimer(1, &CThreadPool::_OnPollTimer);
		const STimerUID& sIncreaseSizeTimerUID = m_oTimerMgr.StartTimer(5, &CThreadPool::_OnIncreaseSizeTimer);
		const STimerUID& sDecreaseSizeTimerUID = m_oTimerMgr.StartTimer(1000, &CThreadPool::_OnDecreaseSizeTimer);
		while (!m_bManagerThreadStoped)
		{
			auto oCurTime = std::chrono::steady_clock::now();
			UINT64 qwPrevTime = m_qwCurTime;
			m_qwCurTime = (UINT64)((oCurTime - m_oStartTime).count() / 1000000);
			UINT32 dwPollCount = m_oTimerMgr.Poll(m_qwCurTime - qwPrevTime);
			for (UINT32 i = 0; i != 10 && !m_oManagerTasks.IsEmpty(); ++i)
			{
				++dwPollCount;
				const std::function<void()>& cbTask = m_oManagerTasks.Front();
				cbTask();
				m_oManagerTasks.Pop();
			}
			int dwTimeout = dwPollCount == 0 ? (int)(m_oTimerMgr.GetNextPollTime() - m_oTimerMgr.GetCurTime()) : 0;
#if _WIN32
			(void)m_oManagerTaskSemaphore.try_acquire_for(std::chrono::milliseconds(dwTimeout));
#else
			int nGot = epoll_wait(m_nEpollFD, m_aEpollEvents, MAX_EPOLL_EVENT_COUNT, dwTimeout);
			_HandleEpollEvent(nGot);
#endif
		}
#if !_WIN32
		DelFD(m_nEventFD);
		ReleasePollEventCallback(qwEventCallback);
#endif
		m_oTimerMgr.StopTimer(sDecreaseSizeTimerUID);
		m_oTimerMgr.StopTimer(sIncreaseSizeTimerUID);
		m_oTimerMgr.StopTimer(sPollTimerUID);
		for (auto&& item : m_oThreadPools)
		{
			DestroyThreadPool(item.first);
		}
		for (auto&& item : m_oThreads)
		{
			if (item.second->joinable())
				item.second->join();
		}
		m_poManagerThread = nullptr;
	}

	void CThreadPool::_PostManagerTask(const std::function<void()>& cbFunction)
	{
		m_oManagerTasks.Push(cbFunction);
		if (!m_oIsManagerTaskQueueNotEmpty.test_and_set(std::memory_order_release))
		{
#if _WIN32
			m_oManagerTaskSemaphore.release();
#else
			UINT64 v = 1;
			write(m_nEventFD, &v, sizeof(v));
#endif
		}
	}

	void CThreadPool::_OnPollTimer(const STimerUID&, UINT32, UINT32)
	{
		for (auto it = m_oPollPerMilliSecond.begin(); it != m_oPollPerMilliSecond.end();)
		{
			auto itCur = it++;
			SPollerContainer& rsPollerContainer = **itCur;
			if (rsPollerContainer.m_sPoller.IsStoped())
			{
				m_oPollPerMilliSecond.erase(itCur);
				continue;
			}
			rsPollerContainer.m_cbOnPoll();
		}
	}

	void CThreadPool::_OnIncreaseSizeTimer(const STimerUID&, UINT32, UINT32)
	{
		for (auto it = m_oThreadPools.begin(); it != m_oThreadPools.end(); ++it)
		{
			auto poThreadPool = it->first;
			SThreadPoolInfo& roInfo = it->second;
			if (poThreadPool->m_bToBeDestroyed)
				continue;
			auto dwThreadCount = poThreadPool->m_dwThreadCount;
			auto dwMaxThreadCount = poThreadPool->m_dwMaxThreadCount;
			auto dwMinThreadCount = poThreadPool->m_dwMinThreadCount;
			auto poFront = poThreadPool->m_oTasks.Front();
			if (dwThreadCount < dwMaxThreadCount &&
				((roInfo.m_pcbTaskHead != nullptr && poFront == roInfo.m_pcbTaskHead) || dwThreadCount < dwMinThreadCount))
			{
				++poThreadPool->m_dwThreadCount;
				++poThreadPool->m_oIdleThreadCount;
				auto poThread = std::make_shared<std::thread>([poThreadPool]() 
				{
					poThreadPool->_ThreadRun();
				});
				m_oThreads[poThread->get_id()] = poThread;
			}
			roInfo.m_pcbTaskHead = poFront;
		}
	}

	void CThreadPool::_OnDecreaseSizeTimer(const STimerUID&, UINT32 dwCurrent, UINT32)
	{
		UINT32 dwMask = 7;
		bool bCheck = (dwCurrent & dwMask) == 0;
		for (auto it = m_oThreadPools.begin(); it != m_oThreadPools.end(); ++it)
		{
			auto poThreadPool = it->first;
			SThreadPoolInfo& roInfo = it->second;
			if (poThreadPool->m_bToBeDestroyed)
			{
				m_oToRemovedThreads.push_back(poThreadPool);
				continue;
			}
			roInfo.m_dwIdleThreadCountPerSecond += poThreadPool->m_oIdleThreadCount;
			if (bCheck)
			{
				if ((UINT32)roInfo.m_dwIdleThreadCountPerSecond > dwMask + 1 &&
					poThreadPool->m_dwThreadCount > poThreadPool->m_dwMinThreadCount)
				{
					--poThreadPool->m_dwThreadCount;
					poThreadPool->m_oTasks.Push(nullptr);
				}
				roInfo.m_dwIdleThreadCountPerSecond = 0;
			}
		}
		for (std::size_t i = 0; i != m_oToRemovedThreads.size(); ++i)
			m_oThreadPools.erase(m_oToRemovedThreads[i]);
		m_oToRemovedThreads.clear();
	}

	void CThreadPool::Run(const std::function<void()>& cbFunction)
	{
		if (!cbFunction)
			return;
		if (m_bToBeDestroyed)
			return;
		m_oTasks.Push(cbFunction);
		m_poSemaphore->release();
	}

	void CThreadPool::_ThreadRun()
	{
		while (!m_bToBeDestroyed)
		{
			--m_oIdleThreadCount;
			std::function<void()> cbTask;
			bool bBreak = false;
			while (m_oTasks.Pop(cbTask))
			{
				if (!cbTask)
				{
					bBreak = true;
					break;
				}
				cbTask();
			}
			++m_oIdleThreadCount;
			if (bBreak)
				break;
			m_poSemaphore->acquire();
		}
		--m_oIdleThreadCount;
		auto oThreadID = std::this_thread::get_id();
		_PostManagerTask([oThreadID]() 
		{
			m_oThreads.erase(oThreadID);
		});
	}

#if !_WIN32
	UINT64 CThreadPool::CreatePollEventCallback(const DOnPollEvent& cbOnPollEvent)
	{
		CSpinLock oLock(m_oPollEventCallbackSpinLockFlag);
		auto oObj = m_oPollEventCallbackObjMgr.CreateObj(0);
		oObj->m_cbOnPollEvent = cbOnPollEvent;
		return oObj->GetUID();
	}

	void CThreadPool::ReleasePollEventCallback(UINT64 qwPollEventCallbackUID)
	{
		CSpinLock oLock(m_oPollEventCallbackSpinLockFlag);
		m_oPollEventCallbackObjMgr.ReleaseObj(qwPollEventCallbackUID);
	}

	bool CThreadPool::AddFD(int dwFD, UINT32 dwEvents, UINT64 qwPollEventCallbackUID)
	{
		epoll_event ev;
		ev.data.u64 = qwPollEventCallbackUID;
		ev.events = dwEvents;
		if (0 > epoll_ctl(m_nEpollFD, EPOLL_CTL_ADD, dwFD, &ev))
			return false;
		return true;
	}

	bool CThreadPool::ModFD(int dwFD, UINT32 dwEvents, UINT64 qwPollEventCallbackUID)
	{
		epoll_event ev;
		ev.data.u64 = qwPollEventCallbackUID;
		ev.events = dwEvents;
		if (0 > epoll_ctl(m_nEpollFD, EPOLL_CTL_MOD, dwFD, &ev))
			return false;
		return true;
	}

	bool CThreadPool::DelFD(int dwFD)
	{
		if (0 > epoll_ctl(m_nEpollFD, EPOLL_CTL_DEL, dwFD, NULL))
			return false;
		return true;
	}

	void CThreadPool::_HandleEpollEvent(int nGot)
	{
		if (nGot > 0)
		{
			for (int j = 0; j != nGot; ++j)
			{
				DOnPollEvent cbOnPollEvent;
				{
					CSpinLock oLock(m_oPollEventCallbackSpinLockFlag);
					auto poCallbackObj = m_oPollEventCallbackObjMgr.GetObj(m_aEpollEvents[j].data.u64);
					if (poCallbackObj != nullptr)
						cbOnPollEvent = poCallbackObj->m_cbOnPollEvent;
				}
				if (cbOnPollEvent)
					cbOnPollEvent(m_aEpollEvents[j].events);
			}
		}
	}
#endif

	CThreadPool::SPoller& CThreadPool::PollPerMilliSecond(const std::function<void()>& cbPollPerMilliSecond)
	{
		std::list<std::shared_ptr<SPollerContainer>>& roPollPerMilliSecond = m_oPollPerMilliSecond;
		std::shared_ptr<SPollerContainer> poPollerContainer = std::make_shared<SPollerContainer>(cbPollPerMilliSecond);
		_PostManagerTask([poPollerContainer, &roPollPerMilliSecond]()
		{
			if (poPollerContainer->m_sPoller.IsStoped())
				return;
			roPollPerMilliSecond.push_back(poPollerContainer);
		});
		return poPollerContainer->m_sPoller;
	}
}
